/*
Author and other information
*/
#include <linux/kernel.h>
#include <linux/module.h>
#include <linux/fs.h>
#include <linux/cdev.h>
#include <linux/init.h>
#include <linux/netfilter.h>
#include <linux/netfilter_ipv4.h>
#include <linux/ip.h>
#include <linux/skbuff.h>
#include <linux/tcp.h>
#include <linux/udp.h>
#include <linux/in.h>
#include <linux/timer.h>
#include <linux/list.h>
#include <linux/net.h>
#include <linux/spinlock.h>
#include <linux/sched.h> 

#include "nat_traversal.h"


MODULE_AUTHOR("Stefan-Liviu Toma");
MODULE_DESCRIPTION("UDP packet interpreter");
MODULE_LICENSE("GPL");

/* Device driver specific data */
struct NatTraversalDeviceData myNatTraversal;


/* this function sends a message through a socket */
int my_socket_send(struct socket *sock, struct sockaddr_in *addr, unsigned char *buf, int len)
{
    struct msghdr msg;
    struct iovec iov;
    mm_segment_t oldfs;
    int size = 0;

    if (sock->sk==NULL)
        return 0;

    iov.iov_base = buf;
    iov.iov_len = len;

    msg.msg_flags = 0;
    msg.msg_name = addr;
    msg.msg_namelen  = sizeof(struct sockaddr_in);
    msg.msg_control = NULL;
    msg.msg_controllen = 0;
    msg.msg_iov = &iov;
    msg.msg_iovlen = 1;
    msg.msg_control = NULL;

    oldfs = get_fs();
    set_fs(KERNEL_DS);
    size = sock_sendmsg(sock,&msg,len);
    set_fs(oldfs);

    return size;
}


/* This function sends a PacketData structure to a destination */
static void sendDataToClient(int realPort, int realIP, 
                             int responseCode, PacketData_t *packetData)
{
    PacketData_t *respPacketData;
    //-----------------------
    struct socket *sock_send;
    struct sockaddr_in addr_send;
    int size, err;
    char *buffer;

    
    /* Trimit datele cilentID catre rP si rIP si codul respCode*/
    /*---------------Trimitere efectiva -----------------------*/
    respPacketData = kzalloc(sizeof( * respPacketData), GFP_ATOMIC);
    if (!respPacketData)
        return;
    //construim packetData
    respPacketData->myPort = 0;
    respPacketData->myIP = 0;
    respPacketData->myID = 0;
    respPacketData->command = 0;
    respPacketData->cmdArgs = 0;
    respPacketData->responseCode = responseCode;
    if (responseCode == CLIENT_FOUND)
    {
        respPacketData->response[0] = packetData->response[0];
        respPacketData->response[1] = packetData->response[1];
        respPacketData->response[2] = packetData->response[2];
        respPacketData->response[3] = packetData->response[3];
    }
    printk("KERN:SDC: retcode = %i, rsp->retCode = %i \n", responseCode, respPacketData->responseCode);
    printk("KERN:Am primit: %i %i %i %i %i %i %i %i %i %i\n",
        packetData->myPort, packetData->myIP, packetData->myID, packetData->command,
        packetData->cmdArgs, packetData->responseCode, packetData->response[0],
        packetData->response[1], packetData->response[2], packetData->response[3]);
    //End of construire packet

    //partea de trimitere
    sock_send = kmalloc(sizeof( *sock_send), GFP_ATOMIC);
    if (!sock_send)
    {
        printk("Eroare alocare sock_send\n");
        return;
    }

    if ( (err = sock_create(AF_INET, SOCK_DGRAM, IPPROTO_UDP, &sock_send)) < 0 )
    {
        printk(": Could not create a datagram socket, error = %d\n", -ENXIO);
        return;
    }

    memset(&addr_send, 0, sizeof(struct sockaddr));

    addr_send.sin_family = AF_INET;
    addr_send.sin_addr.s_addr = realIP;
    addr_send.sin_port = htons(realPort);


    buffer = kmalloc( sizeof( *respPacketData), GFP_ATOMIC);
    if (!buffer)
    {
        sock_release(sock_send);
        return;
    }
    memset(buffer, 0, sizeof( *respPacketData));
    memcpy(buffer, respPacketData, sizeof( *respPacketData));
    my_socket_send(sock_send, &addr_send, buffer, sizeof( *respPacketData));

    sock_release(sock_send);
    kfree(packetData);

    /*---------------End of trimitere efectiva ----------------*/

}


/* This is the handler used by scheduled work structure */
void my_work_handler(struct work_struct *work)
{
    struct NatTraversalDeviceData *my_data = container_of(work, struct NatTraversalDeviceData, my_work);
    //printk(" Intra work handler\n");
    printk(" WKH rp = %i, rip = %i, ret = %i \n" ,
        my_data->details->realP, my_data->details->realIP, my_data->details->retCode);
    sendDataToClient(my_data->details->realP, my_data->details->realIP,
        my_data->details->retCode, my_data->details->packetData);
} 



//obtain client details
PacketData_t *GetClientEndpoints(int clientID, int *ret)
{
    PacketData_t *packetData;
    struct list_head *iterator, *n;
    ClientData_t *client;

    if (!clientExists(clientID))
    {
        *ret = GENERAL_ERROR;
        return NULL;
    }

    packetData = kzalloc( sizeof( *packetData), GFP_ATOMIC);
    if (!packetData)
    {
        *ret = GENERAL_ERROR;
        return NULL;
    }

    list_for_each_safe(iterator, n, &myNatTraversal.allClients)
    {
        client = list_entry(iterator, struct ClientData, list);
        if (client->clientID == clientID)
        {
            packetData->responseCode = CLIENT_FOUND;
            packetData->response[0] = client->privatePort;
            packetData->response[1] = client->privateIP;
            packetData->response[2] = client->publicPort;
            packetData->response[3] = client->publicIP;
            *ret = CLIENT_FOUND;
        }
    }
    if ( *ret != CLIENT_FOUND )
        *ret = CLIENT_NOT_FOUND;

    return packetData;

}


static void processPacketData(int realPort, int realIP, PacketData_t *packetData)
{
    int ret;
    int rp = 0, rip = 0;
    int retVal = 0;
    PacketData_t *newPacketData;
    rp = realPort;
    rip = realIP;

    //printk("comanda = %i\n", packetData->command);

    switch(packetData->command)
    {
    case REGISTER:
        if (myNatTraversal.clientNo < MAX_NO_CLIENTS )
        {
            ret = addClient(realPort, realIP, packetData);

            myNatTraversal.details->realP = realPort;
            myNatTraversal.details->realIP = realIP;
            myNatTraversal.details->retCode = ret;
            myNatTraversal.details->packetData = packetData;
            INIT_WORK(&myNatTraversal.my_work, my_work_handler);
            schedule_work(&myNatTraversal.my_work); 
        }
        else
        {
            myNatTraversal.details->realP = realPort;
            myNatTraversal.details->realIP = realIP;
            myNatTraversal.details->retCode = SERVER_FULL;
            myNatTraversal.details->packetData = packetData;
            INIT_WORK(&myNatTraversal.my_work, my_work_handler);
            schedule_work(&myNatTraversal.my_work);
        }
        break;

    case UNREGISTER:
        ret = deleteClient(packetData->myID);
        myNatTraversal.details->realP = realPort;
        myNatTraversal.details->realIP = realIP;
        myNatTraversal.details->retCode = ret;
        myNatTraversal.details->packetData = packetData;
        INIT_WORK(&myNatTraversal.my_work, my_work_handler);
        schedule_work(&myNatTraversal.my_work);
        break;

    case REQUEST_DATA:

        //daca nu este inregistrat, semnalam acest lucru
        if (!clientExists(packetData->myID))
        {
            myNatTraversal.details->realP = realPort;
            myNatTraversal.details->realIP = realIP;
            myNatTraversal.details->retCode = NOT_REGISTERED;
            myNatTraversal.details->packetData = packetData;
            INIT_WORK(&myNatTraversal.my_work, my_work_handler);
            schedule_work(&myNatTraversal.my_work);
        }
        else
        {
            //obtinem date despre clientul cerut
            newPacketData = GetClientEndpoints(packetData->cmdArgs, &retVal);
            if (!newPacketData)
                myNatTraversal.details->packetData = packetData;
            else
            {
                myNatTraversal.details->packetData = newPacketData;
                //kfree(packetData); //daca nu, nu se mai elibereaza deloc
            }
            myNatTraversal.details->realP = realPort;
            myNatTraversal.details->realIP = realIP;
            myNatTraversal.details->retCode = retVal;
            INIT_WORK(&myNatTraversal.my_work, my_work_handler);
            schedule_work(&myNatTraversal.my_work);
        }

        break;

    default:
        printk(" Comanda nerecunoscuta\n");
        myNatTraversal.details->realP = realPort;
        myNatTraversal.details->realIP = realIP;
        myNatTraversal.details->retCode = INVALID_COMMAND;
        myNatTraversal.details->packetData = packetData;
        INIT_WORK(&myNatTraversal.my_work, my_work_handler);
        schedule_work(&myNatTraversal.my_work);
        break;

    }
}



static unsigned int hook_function (unsigned int hooknum,
struct sk_buff *skb,
    const struct net_device *in,
    const struct net_device *out,
    int (*okfn)(struct sk_buff *))
{
    struct iphdr *iph = ip_hdr(skb);
    struct udphdr *udph;
    unsigned char *effectiveData;
    PacketData_t *packetData;
    int IPSOURCE = iph->saddr;

    /* process packet if server is enabled*/
    if( !myNatTraversal.enabled )
        return NF_ACCEPT;


    /* We don't bother for NULL pointers */
    if ( !iph )
        return NF_ACCEPT;
    /* We analyze UDP packets*/
    if( iph->protocol == IPPROTO_UDP )
    {
        /* Show packet info: */
        printk( " Pachet intrat: L = %i, PROT = %i, SA = %x, DA = %x \n",
            iph->ihl, iph->protocol, iph->saddr, iph->daddr);

        //printk(" In prima parte: IPSOURCE = %i. addr = %i.", IPSOURCE, iph->saddr);
        udph = (struct udphdr *)( skb->data + iph->ihl * 4);
        if( !udph )
            printk( "UDPheader NULL \n");
        else
        {
            if( ntohs(udph->dest) != SERVER_PORT )
            {
                //printk(" Alt port decat al serverului: %i\n", ntohs(udph->dest));
                return NF_ACCEPT;
            }
            /*printk( " Detalii UDP 2: S = %i, D = %i, len = %i, crc = %x \n",
                ntohs(udph->source), ntohs(udph->dest), ntohs(udph->len), udph->check);*/

            effectiveData = skb->data + iph->ihl * 4 + 8;

            packetData = getPacketData(effectiveData);
            if (!packetData)
                printk("Wasn't able to get packetData\n");
            else
            {
                //printk(" In hook ultima parte: IPSOURCE = %i. addr = %i.", IPSOURCE, iph->saddr);
                processPacketData(ntohs(udph->source), IPSOURCE, packetData);
            }
        }
    }

    return NF_ACCEPT;
}

/**/

PacketData_t *getPacketData(unsigned char *data)
{
    PacketData_t *packetData;

    packetData = kmalloc(sizeof( * packetData), GFP_ATOMIC);
    if (!packetData)
        return NULL;

    memcpy(packetData, data, sizeof(*packetData));

    printk("Detalii date: %x, %x, %x, %x, %x .\n",
        packetData->myPort, packetData->myIP, packetData->myID, packetData->command, packetData->cmdArgs);

    return packetData;
}


int clientExists(int clientID)
{
    struct list_head *iterator, *n;
    ClientData_t *client;

    list_for_each_safe(iterator, n, &myNatTraversal.allClients)
    {
        client = list_entry(iterator, struct ClientData, list);

        if (client->clientID == clientID)
            return 1;
    }

    return 0;
}



int addClient(int publicPort, int publicIP, PacketData_t *packetData)
{
    ClientData_t *clientData;

    clientData = kmalloc(sizeof( * clientData), GFP_ATOMIC);
    if (!clientData)
        return GENERAL_ERROR;
    if (packetData->myID == 0)
        return INVALID_ID;

    if( !clientExists(packetData->myID) )
    {
        clientData->publicPort = publicPort;
        clientData->publicIP = publicIP;
        clientData->privatePort = packetData->myPort;
        clientData->privateIP = packetData->myIP;
        clientData->clientID = packetData->myID;
        spin_lock(&myNatTraversal.mylock);
        myNatTraversal.clientNo ++;
        INIT_LIST_HEAD(&clientData->list);
        printk(" Adaug client = %i\n", clientData->clientID);
        list_add(&clientData->list, &myNatTraversal.allClients);
        spin_unlock(&myNatTraversal.mylock);
    }
    else
        return INVALID_ID;

    return REGISTER_OK;

}


void addClientFromUser(ClientData_t *clientData)
{
    if( !clientExists(clientData->clientID) )
    {
        spin_lock(&myNatTraversal.mylock);
        myNatTraversal.clientNo ++;
        INIT_LIST_HEAD(&clientData->list);
        printk(" Adaug client = %i\n", clientData->clientID);
        list_add(&clientData->list, &myNatTraversal.allClients);
        spin_unlock(&myNatTraversal.mylock);
    }

}



int deleteClient(int clientID)
{
    struct list_head *iterator, *n;
    ClientData_t *client;

    printk("Stergem client = %i\n", clientID);

    if (clientID == 0)
        return INVALID_ID;

    list_for_each_safe(iterator, n, &myNatTraversal.allClients)
    {
        client = list_entry(iterator, struct ClientData, list);

        if (client->clientID == clientID)
        {
            spin_lock(&myNatTraversal.mylock);
            list_del(iterator);
            myNatTraversal.clientNo --;
            kfree(client);
            spin_unlock(&myNatTraversal.mylock);
            return UNREGISTER_OK;
        }
    }

    return UNREGISTER_NOT_OK;



}


void deleteAllClients()
{
    struct list_head *iterator, *n;
    ClientData_t *client;

    list_for_each_safe(iterator, n, &myNatTraversal.allClients)
    {
        client = list_entry(iterator, struct ClientData, list);
        deleteClient(client->clientID);
    }

}



static int my_open(struct inode *inode, struct file *file)
{
    return 0;
}

static int my_close(struct inode *inode, struct file *file)
{
    return 0;
}

//functia de control a device-ului
static int my_ioctl(struct inode *inode, struct file *file,unsigned int cmd, unsigned long arg)
{
    void *buffer;
    int i;
    struct list_head *iterator, *n;
    //PacketData_t *packetData;
    ClientData_t *client;
    ClientData_t *newClient;
    int clientID = 0;

    switch(cmd)
    {

    case IOCTL_SERVER_ENABLE:
        myNatTraversal.enabled = 1;
        break;

    case IOCTL_SERVER_DISABLE:
        myNatTraversal.enabled = 0;
        break;

    case IOCTL_SERVER_ADD_CLIENT:
        newClient = kmalloc(sizeof( *newClient), GFP_ATOMIC);
        if(!newClient)
            return -ENOMEM;
        if( copy_from_user(newClient, (void *) arg, sizeof( *newClient)) )
            return -EFAULT;
        addClientFromUser(newClient);
        break;

    case IOCTL_SERVER_DEL_CLIENT:
        if( copy_from_user(&clientID,(int *) arg, sizeof(int)) )
            return -EFAULT;
        if (!clientExists(clientID))
            return -EINVAL;
        else
            deleteClient(clientID);
        break;

    case IOCTL_SERVER_GET_CLIENTS:
        buffer = kmalloc(myNatTraversal.clientNo * sizeof( *client), GFP_ATOMIC);
        if(!buffer)
            return -ENOMEM;
        i = 0;
        spin_lock(&myNatTraversal.mylock);
        list_for_each_safe(iterator, n, &myNatTraversal.allClients)
        {
            client = list_entry(iterator, struct ClientData, list);
            /* Just the client information*/
            memcpy(buffer + i * sizeof( *client), client, 5 * sizeof(int));
            i++;
        }
        spin_unlock(&myNatTraversal.mylock);
        if( (void *)arg == NULL )
        {
            return myNatTraversal.clientNo;
        }

        if( copy_to_user( (void *)arg, buffer, myNatTraversal.clientNo * 5 * sizeof(int)) )
        {
            kfree(buffer);
            return -EFAULT;
        }

        kfree(buffer);
        return myNatTraversal.clientNo;

        break;
    }

    return 0;
}


struct file_operations my_fops = {
    .owner 		= THIS_MODULE,
    .open 		= my_open,
    .release 	= my_close,
    .ioctl 		= my_ioctl
};


int __init hook_init (void)
{
    int err;

    myNatTraversal.details = kmalloc( sizeof(struct Details), GFP_ATOMIC);
    if (!myNatTraversal.details)
    {
        printk("Cannot alloc details \n");
        return -ENOMEM;
    }
    //inregistrare regiune chardev
    err = register_chrdev_region(MKDEV(NAT_TRAVERSAL_MAJOR, 0), 1,"NatTraversal");
    if (err != 0) 
        return err;//we cannot register the driver

    cdev_init(&myNatTraversal.cdev, &my_fops);
    cdev_add(&myNatTraversal.cdev, MKDEV(NAT_TRAVERSAL_MAJOR, 0), 1);

    //Netfilter specific init
    myNatTraversal.hook_in.owner		= THIS_MODULE;
    myNatTraversal.hook_in.hook			= hook_function;
    myNatTraversal.hook_in.hooknum		= NF_IP_LOCAL_IN;
    myNatTraversal.hook_in.pf			= PF_INET;
    myNatTraversal.hook_in.priority		= NF_IP_PRI_FIRST;

    /* Extra initialization*/
    INIT_LIST_HEAD(&myNatTraversal.allClients);
    spin_lock_init(&myNatTraversal.mylock);
    myNatTraversal.clientNo = 0;
    //temporar
    myNatTraversal.enabled = 1;

    nf_register_hook(&myNatTraversal.hook_in);
    printk("NAT Traversal server started.\n");
    return 0;
}

void __exit hook_exit (void)
{
    //Delete all clients
    deleteAllClients();

    //Unregister hooks
    nf_unregister_hook(&myNatTraversal.hook_in);

    //Unregister chardev region
    cdev_del(&myNatTraversal.cdev);
    unregister_chrdev_region(MKDEV(NAT_TRAVERSAL_MAJOR, 0), 1);
    printk("NAT Traversal server unloaded.\n");
}

module_init (hook_init);
module_exit (hook_exit);

